import ast
import math
import re
from decimal import Decimal, getcontext

from lib.CountFunc import *

# 配置高精度计算上下文
context = getcontext()
context.prec = 1000  # 1000位精度
context.Emin = -1000000  # 最小指数
context.Emax = 1000000  # 最大指数


class CountModel:
    """支持高精度计算的数值模型（含复数支持）"""

    def __init__(self, number):
        # 处理复数输入
        if isinstance(number, complex):
            self.real = Decimal(str(number.real))
            self.imag = Decimal(str(number.imag))
        elif isinstance(number, str):
            # 尝试使用 evaluate_expression 计算表达式的值
            result = self.evaluate_expression(number)
            if result is not None:
                num = result.get_original_number()
                self.real = Decimal(str(num.real))
                self.imag = Decimal(str(num.imag))
            else:
                raise ValueError(f"无法解析输入: {number}")
        else:
            try:
                num = complex(number)  # 统一转换为复数处理
                self.real = Decimal(str(num.real))
                self.imag = Decimal(str(num.imag))
            except:
                raise ValueError(f"无法解析输入: {number}")

    def get_original_number(self):
        """获取原始数值（复数形式）"""
        return complex(float(self.real), float(self.imag))

    def scientificCounting(self, number: Decimal):
        """科学计数法表示"""
        # 将 Decimal 转换为科学计数法字符串
        scientific_str = '{:.10e}'.format(number)
        # 分割尾数和指数部分
        mantissa_str, exponent_str = scientific_str.split('e')
        mantissa = Decimal(mantissa_str)
        exponent = int(exponent_str)
        return f"{mantissa} * 10^{exponent}"

    def __repr__(self):
        """字符串格式化（处理复数显示）"""
        if self.imag == 0:
            return f"{self.real:.10f}"  # 实部显示
        elif self.real == 0:
            return f"{self.imag:.10f}i"  # 纯虚数
        else:
            return f"{self.real:.10f}{self.imag:.10f}"

    def __str__(self):
        """字符串格式化（处理复数显示）"""
        if self.imag == 0:
            return self.scientificCounting(self.real)
            # f"{self.real:.10f}"  # 实部显示
        elif self.real == 0:
            return f"{self.scientificCounting(self.imag)} i"
            # f"{self.imag:.10f}i"  # 纯虚数
        else:
            return f"{self.scientificCounting(self.real)}{'+' if self.imag > 0 else ''}{self.scientificCounting(self.imag)} i"
            # f"{self.real:.10f}{'+' if self.imag > 0 else ''}{self.imag:.10f}i"

    @staticmethod
    def evaluate_expression(expression):
        """安全计算含math函数的表达式（支持阶乘!和绝对值|x|）"""
        # 预处理1：将|x|替换为abs(x)（支持绝对值运算）
        expression = re.sub(r'\|([^|]+)\|', r'abs(\1)', expression)  # 匹配两个竖线之间的任意内容

        # 预处理2：将x!替换为math.gamma(x + 1)（支持阶乘）
        expression = re.sub(r'(\([^!]+\)|\w+|\))!', r'math.gamma(\1 + 1)', expression)

        allowed = {
            'sin': math.sin, 'cos': math.cos, 'tan': math.tan,
            'sqrt': math.sqrt, 'exp': math.exp, 'log': math.log,
            'pi': math.pi, 'e': math.e, 'i': 1j,
            'math': math, 'gamma': math.gamma,
            'abs': abs  # 显式允许abs函数（处理绝对值）
        }
        # 若CountFunc内存在函数插件,将其添加到allowed中
        for func_name, func in CountFunc.__dict__.items():
            if callable(func):
                allowed[func_name] = func
        # print(f"CountModel:63 <UNK>: {expression}")
        # print("CountModel:64", allowed)
        try:
            tree = ast.parse(expression, mode='eval')
            # print(tree.body.__dict__)
            for node in ast.walk(tree):
                # print(node, ast.Name, node.id)
                if isinstance(node, ast.Name) and node.id not in allowed:
                    raise ValueError(f"不允许的函数/变量: {node.id}")
            code = compile(tree, '<string>', 'eval')
            # print(code.__dict__)
            result = eval(code, {'__builtins__': {}}, allowed)
            # print(result)
            return CountModel(result)
        except Exception as e:
            print(f"计算错误: {e}")
            return None

    # 运算符重载（加减乘除）
    def __add__(self, other):
        other = other.get_original_number() if isinstance(other, CountModel) else complex(other)
        return CountModel(self.get_original_number() + other)

    def __sub__(self, other):
        other = other.get_original_number() if isinstance(other, CountModel) else complex(other)
        return CountModel(self.get_original_number() - other)

    def __mul__(self, other):
        other = other.get_original_number() if isinstance(other, CountModel) else complex(other)
        return CountModel(self.get_original_number() * other)

    def __truediv__(self, other):
        other = other.get_original_number() if isinstance(other, CountModel) else complex(other)
        if other == 0:
            raise ZeroDivisionError("除以零错误")
        return CountModel(self.get_original_number() / other)
